// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 *
 * Description: End-to-End Euler Net Boost sample.
 */
#include <uapi/linux/in.h>
#include <uapi/linux/if_ether.h>
#include <linux/if_vlan.h>
#include <uapi/linux/ip.h>
#include <uapi/linux/tcp.h>
#include <uapi/linux/bpf.h>

#include <bpf/bpf_endian.h>
#include <bpf/bpf_helpers.h>

#define IP_MF           0x2000
#define IP_OFFSET       0x1FFF
#define CSUM_SHIFT_BITS	16

#define SOCKOPS_SUCC	1
#define SOCKOPS_FAIL	0

static long g_rx_dst;

static struct {
	u8 eth_laddr[ETH_ALEN];
	u8 eth_raddr[ETH_ALEN];
	u8 updated;
} g_eth_addr;

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(key_size, sizeof(u32));
	__uint(value_size, sizeof(u32));
	__uint(max_entries, 1);
} speed_port SEC(".maps");

static inline bool is_speed_flow(u32 port)
{
	u32 key = 0;
	u32 *val = bpf_map_lookup_elem(&speed_port, &key);

	return val != NULL && *val == port;
}

SEC("enboost_sockops")
int enboost_sockops_prog(struct bpf_sock_ops *skops)
{
	long dst;

	if (!is_speed_flow(bpf_ntohl(skops->remote_port)))
		return SOCKOPS_SUCC;

	switch (skops->op) {
	case BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB:
		dst = bpf_get_sk_rx_dst(skops);
		if (IS_ERR((void *)dst))
			return SOCKOPS_FAIL;

		g_rx_dst = dst;
		bpf_sock_ops_cb_flags_set(skops, BPF_SOCK_OPS_STATE_CB_FLAG);
		break;

	case BPF_SOCK_OPS_STATE_CB:
		if (skops->args[0] == BPF_TCP_LAST_ACK ||
		    skops->args[1] == BPF_TCP_CLOSE) {
			g_rx_dst = 0;
			g_eth_addr.updated = 0;
			bpf_sock_ops_cb_flags_set(skops, skops->bpf_sock_ops_cb_flags &
						  ~BPF_SOCK_OPS_STATE_CB_FLAG);
		}
		break;
	}

	return SOCKOPS_SUCC;
}

SEC("enboost_ingress")
int enboost_ingress_prog(struct xdp_md *ctx)
{
	void *data = (void *)(long)ctx->data;
	void *data_end = (void *)(long)ctx->data_end;
	char ingress_dev[IFNAMSIZ] = "eth0";
	struct ethhdr *ehdr;
	struct tcphdr *thdr;
	struct iphdr *ihdr;

	ehdr = (struct ethhdr *)data;
	if (ehdr + 1 > data_end)
		return XDP_PASS;

	if (eth_type_vlan(ehdr->h_proto)) {
		struct vlan_hdr *vhdr;
		__be16 h_proto;

		vhdr = (struct vlan_hdr *)(ehdr + 1);
		if (vhdr + 1 > data_end)
			return XDP_PASS;

		h_proto = vhdr->h_vlan_encapsulated_proto;
		if (h_proto != bpf_htons(ETH_P_IP))
			return XDP_PASS;

		ihdr = (struct iphdr *)(vhdr + 1);
	} else {
		if (ehdr->h_proto != bpf_htons(ETH_P_IP))
			return XDP_PASS;

		ihdr = (struct iphdr *)(ehdr + 1);
	}

	if (ihdr + 1 > data_end)
		return XDP_PASS;

	if (ihdr->ihl != 5 || ihdr->protocol != IPPROTO_TCP)
		return XDP_PASS;

	if (ihdr->frag_off & bpf_htons(IP_MF | IP_OFFSET))
		return XDP_PASS;

	thdr = (struct tcphdr *)(ihdr + 1);
	if (thdr + 1 > data_end)
		return XDP_PASS;

	if (thdr->syn || thdr->fin || thdr->rst)
		return XDP_PASS;

	if (!is_speed_flow(bpf_ntohs(thdr->source)))
		return XDP_PASS;

	if (!g_eth_addr.updated) {
		__builtin_memcpy(g_eth_addr.eth_laddr, ehdr->h_dest, ETH_ALEN);
		__builtin_memcpy(g_eth_addr.eth_raddr, ehdr->h_source, ETH_ALEN);
		g_eth_addr.updated = 1;
	}

	if (g_rx_dst)
		bpf_xdp_set_skb_dst(ctx, g_rx_dst);

	/* Whether to use it based on the actual scenario */
	bpf_change_skb_dev(ctx, ingress_dev);
	return XDP_BOOST;
}

static inline void ipv4_csum(struct iphdr *ihdr)
{
	u32 csum = 0;
	u16 *next_ip_u16 = (u16 *)ihdr;

	ihdr->check = 0;
	for (size_t i = 0; i < sizeof(struct iphdr) >> 1; i++)
		csum += *next_ip_u16++;

	ihdr->check = ~((csum & 0xffff) + (csum >> CSUM_SHIFT_BITS));
}

SEC("enboost_egress")
int enboost_egress_prog(struct __sk_buff *skb)
{
	void *data = (void *)(long)skb->data;
	void *data_end = (void *)(long)skb->data_end;
	char egress_dev[IFNAMSIZ] = "eth0";
	struct ethhdr *ehdr;
	struct tcphdr *thdr;
	struct iphdr *ihdr;
	int ret;

	if (!is_speed_flow(bpf_ntohl(skb->remote_port)))
		return ENET_PASS;

	ihdr = (struct iphdr *)data;
	if (ihdr + 1 > data_end)
		return ENET_PASS;

	if (ihdr->ihl != 5 || ihdr->protocol != IPPROTO_TCP)
		return ENET_PASS;

	thdr = (struct tcphdr *)(ihdr + 1);
	if (thdr + 1 > data_end)
		return ENET_PASS;

	if (thdr->syn || thdr->fin || thdr->rst)
		return ENET_PASS;

	if (!g_eth_addr.updated)
		goto boost;

	ihdr->tot_len = bpf_htons(skb->len);
	ipv4_csum(ihdr);

	ret = bpf_skb_change_head(skb, ETH_HLEN, 0);
	if (ret < 0)
		goto boost;

	data = (void *)(long)skb->data;
	data_end = (void *)(long)skb->data_end;

	ehdr = (struct ethhdr *)data;
	if (ehdr + 1 > data_end)
		return ENET_DROP;

	ehdr->h_proto = bpf_htons(ETH_P_IP);
	__builtin_memcpy(ehdr->h_source, g_eth_addr.eth_laddr, ETH_ALEN);
	__builtin_memcpy(ehdr->h_dest, g_eth_addr.eth_raddr, ETH_ALEN);
boost:
	/* Whether to use it based on the actual scenario */
	bpf_change_skb_dev(skb, egress_dev);
	return ENET_BOOST;
}

char _license[] SEC("license") = "GPL";
